import json

import pymysql

import MySQL
import xml.sax

paper_tags = ('article', 'inproceedings', 'proceedings', 'book', 'phdthesis', 'mastersthesis', 'www')


default_db = 'dblp'
default_table = 'dblp_xml'
error = 0
objs = []
dblp_id = 0


class DataHandler(xml.sax.ContentHandler):
    def __init__(self):
        super().__init__()
        self.current = ''
        self.tag = ''
        self.title = ''
        self.mdate = ''
        self.year = 0
        self.key = ''
        self.info = {}

    def startElement(self, tag, attributes):
        self.current = tag
        if tag in paper_tags:
            self.tag = tag
            if 'mdate' in attributes:
                self.mdate = attributes['mdate']
            if 'key' in attributes:
                self.key = attributes['key']

    def endElement(self, tag):
        global objs
        global error
        global dblp_id
        if tag in paper_tags:
            dblp_id = dblp_id + 1
            row = {
                'id': dblp_id,
                'tag': self.tag,
                'title': self.title,
                'mdate': self.mdate,
                'year': self.year,
                'key': self.key,
                'info': json.dumps(self.info)[:4000]
            }
            objs.append(row)
            if dblp_id % 1000 == 0:
                try:
                    MySQL.insert(default_table, objs)
                except pymysql.err.DataError as e:
                    global error
                    error = error + len(objs)
                    print(e)
                except pymysql.err.OperationalError:
                    MySQL.init(default_db, 0, default_charset='utf8mb4')
                    MySQL.insert(default_table, objs)
                except pymysql.err.ProgrammingError:
                    for obj in objs:
                        try:
                            MySQL.save(default_table, obj)
                        except Exception as e:
                            error = error + 1
                            print(e)
                objs = []
            self.tag = ''
            self.title = ''
            self.mdate = ''
            self.year = 0
            self.key = ''
            self.info = {}
        if tag == 'dblp' and len(objs) > 0:
            try:
                MySQL.insert(default_table, objs)
            except pymysql.err.DataError as e:
                error = error + len(objs)
                print(e)
            except pymysql.err.OperationalError:
                MySQL.init(default_db, 0, default_charset='utf8mb4')
                MySQL.insert(default_table, objs)
            except pymysql.err.ProgrammingError:
                for obj in objs:
                    try:
                        MySQL.save(default_table, obj)
                    except Exception as e:
                        error = error + 1
                        print(e)
            objs = []
        self.current = ''

    def characters(self, content):
        if self.current == 'title':
            self.title = content
        if self.current == 'year':
            self.year = content
        if self.current == 'author' or self.current == 'editor' or self.current == 'school':
            if self.current not in self.info.keys():
                self.info[self.current] = []
            self.info[self.current].append(content)


def resolve(file_name):
    parser = xml.sax.make_parser()
    parser.setFeature(xml.sax.handler.feature_namespaces, 0)
    parser.setContentHandler(DataHandler())
    parser.parse(file_name)
    print('error:' + str(error))


if __name__ == '__main__':
    MySQL.init(default_db, 0, default_charset='utf8mb4')
    for p in range(54):
    # for p in range(14, 15):
        name = '../data/dblp/dblp_' + str(p) + '.xml'
        resolve(name)
        print('resolve(' + name + ') finished')
